#coding:utf-8
from idaapi import *
from xml.dom.minidom import Document
'''
获得危险函数名称及对应位置信息
根据一个给定的XML Schema，使用DOM树的形式从空白文件生成一个XML
'''

doc = Document()  #创建DOM文档对象

bookstore = doc.createElement('bookstore') #创建根元素
bookstore.setAttribute('xmlns:xsi',"https://www.w3.org/2001/XMLSchema-instance")#设置命名空间
bookstore.setAttribute('xsi:noNamespaceSchemaLocation','bookstore.xsd')#引用本地XML Schema
doc.appendChild(bookstore)

book = doc.createElement('book')
book.setAttribute('genre','XML')
bookstore.appendChild(book)

title = doc.createElement('title')
title_text = doc.createTextNode('IDAPython危险函数调用检测') #元素内容写入
title.appendChild(title_text)
book.appendChild(title)

print("Loading the dangerous library file...")
danger_funcs = ["IsProcessorFeaturePresent"]  # 需要寻找的函数名，此函数可以确定CPU支持哪些特性，可使用文本（数据库）读取的方式进行更多危险函数的查询
print("Start searching for danger functions,Potential vulnerabilities will be printed below,Please wait...")

#漏洞名称及地址
vulnerabilities = doc.createElement('Vulnerabilities_DangerFunctions')
book.appendChild(vulnerabilities)

price_size = 0
for func in danger_funcs:	
	addr = LocByName( func ) #通过名称获得地址
	if addr != BADADDR:     #如果地址不是错误的
		#找到交叉引用的地址
		cross_refs = CodeRefsTo( addr, 0 )
		print "Cross References to %s" % func 
		print "-------------------------------"
		for ref in cross_refs: 
			author_name = doc.createElement('Vulnerability-name')
			author_addr  = doc.createElement('Vulnerability-addr')
			author_name_text = doc.createTextNode(func)
			author_addr_text  = doc.createTextNode(hex(ref))
			vulnerabilities.appendChild(author_name)
			vulnerabilities.appendChild(author_addr)
			author_name.appendChild(author_name_text)
			author_addr.appendChild(author_addr_text)
			book.appendChild(vulnerabilities)
			price_size += 2
			
			print "%08x" % ref
			SetColor( ref, CIC_ITEM, 0x0000ff)  #设置函数的颜色为红色
			
#这里写漏洞影响大小             
price = doc.createElement('price')
price_text = doc.createTextNode(str(price_size))
price.appendChild(price_text)
book.appendChild(price)

########### 将DOM对象doc写入文件
f = open('Vulnerability_Analysis_Report_DangerFunctions.xml','w')
f.write(doc.toprettyxml(indent = ''))
f.close()

print("Danger functions search ended")








